// Package main go 实现的压测工具
package main

import (
	"context"
	"flag"
	"fmt"
	"gitee.com/zhaobohan/stress-testing/model"
	"gitee.com/zhaobohan/stress-testing/server"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"syscall"
	"time"
)

type array []string

// String 变更格式化输出的规则
func (a *array) String() string {
	return fmt.Sprint(*a)
}

// Set 向数组中添加元素
func (a *array) Set(s string) error {
	*a = append(*a, s)
	return nil
}

var (
	concurrency uint64 = 1       // 并发数
	totalNumber uint64 = 1       // 请求数(单个并发/协程)
	debugStr           = "false" // 是否是debug
	requestURL         = ""      // 压测的url
	path               = ""      // curl文件路径 http接口压测，自定义参数设置
	verify             = ""      // 验证方法
	headers     array            // 自定义头信息传递给服务器
	body               = ""      // HTTP POST方式传送数据
	maxCon             = 1       // 单个连接最大请求数
	code               = 200     // 成功状态码
	http2              = false   // 是否开http2.0
	keepalive          = false   // 是否开启长连接
	cpuNumber          = 1       // CPU 核数
	timeout     int64  = 0       // 接口超时时间
	appTimeout  int64  = 0       // 压测程序最大执行时间，默认不设置
	redirect           = true    // 是否重定向
)

// init读取对应的配置信息，存储到变量当中
func init() {
	flag.Uint64Var(&concurrency, "c", concurrency, "并发数")
	flag.Uint64Var(&totalNumber, "n", totalNumber, "请求数(单个并发/协程)")
	flag.StringVar(&debugStr, "d", debugStr, "调试模式")
	flag.StringVar(&requestURL, "u", requestURL, "压测地址")
	flag.StringVar(&path, "p", path, "curl文件路径")
	flag.StringVar(&verify, "v", verify, "验证方法")
	flag.Var(&headers, "H", "自定义头信息传递给服务器 示例:-H 'Content-Type: application/json'")
	flag.StringVar(&body, "data", body, "HTTP POST方式传送数据")
	flag.IntVar(&maxCon, "m", maxCon, "单个host最大连接数")
	flag.IntVar(&code, "code", code, "请求成功的状态码")
	flag.BoolVar(&http2, "http2", http2, "是否开 http2.0")
	flag.BoolVar(&keepalive, "k", keepalive, "是否开启长连接")
	flag.Int64Var(&timeout, "t", timeout, "接口超时时间 单位 秒，默认为 30s")
	flag.IntVar(&cpuNumber, "cpuNumber", cpuNumber, "CPU 核数，默认为一核")
	flag.Int64Var(&appTimeout, "timeout", appTimeout, "压测程序最大执行时间 单位 秒,默认一直压测")
	flag.BoolVar(&redirect, "redirect", redirect, "是否重定向")
	flag.Parse()
}

// main go 实现的压测工具
func main() {
	// 用于设置程序运行时可以使用的最大 CPU 核心数，只是可以，不是绝对的，取决于操作系统的调度策略
	runtime.GOMAXPROCS(cpuNumber)

	// 检查参数
	if concurrency == 0 || totalNumber == 0 || (requestURL == "" && path == "") {
		fmt.Printf("示例: go run main.go -c 1 -n 1 -u https://zhaobohan.blog.csdn.net/ \n")
		fmt.Printf("压测地址或curl路径必填 \n")
		fmt.Printf("当前请求参数: -c %d -n %d -d %v -u %s \n", concurrency, totalNumber, debugStr, requestURL)
		flag.Usage()
		return
	}
	debug := strings.ToLower(debugStr) == "true"
	request, err := model.NewRequest(requestURL, verify, code, time.Duration(timeout)*time.Second, debug, path, headers,
		body, maxCon, http2, keepalive, redirect)
	if err != nil {
		fmt.Printf("参数不合法 %v \n", err)
		return
	}
	fmt.Printf("\n 开始启动  并发数:%d 请求数:%d 请求参数: \n", concurrency, totalNumber)
	request.Print()

	// 开始处理
	ctx := context.Background()
	if appTimeout > 0 {
		var cancel context.CancelFunc
		ctx, cancel = context.WithTimeout(ctx, time.Duration(appTimeout)*time.Second)
		defer cancel()
		deadline, ok := ctx.Deadline()
		if ok {
			fmt.Printf(" deadline %s", deadline)
		}
	}

	// 处理 ctrl+c 信号
	ctx, cancelFunc := context.WithCancel(ctx)
	go func() {
		c := make(chan os.Signal)
		signal.Notify(c, syscall.SIGINT)
		<-c
		cancelFunc()
	}()

	// 模拟的是concurrency个用户，每个用户执行totalNumber次请求
	server.Dispose(ctx, concurrency, totalNumber, request)
	return
}
